import pandas as pd
from gensim.models import TfidfModel
from gensim.models.ldamodel import LdaModel
from gensim.corpora import Dictionary

import jieba


def m_head(tmpstr):
    return tmpstr[:1]


def m_mid(tmpstr):
    return tmpstr.find("回 ")


def cut(content):
    s = pd.read_table("./停用词.txt", names=['txt'], engine="python")
    res = []
    for c in jieba.cut(content):
        if (c not in s) and len(c) > 1:
            res.append(c)
    return res

def chapter():
    raw = pd.read_csv("./金庸-射雕英雄传txt精校版.txt", names=['txt'], sep='aaa', encoding="GBK", engine='python')
    raw['head'] = raw.txt.apply(m_head)
    raw['mid'] = raw.txt.apply(m_mid)
    raw['len'] = raw.txt.apply(len)
    chapnum = 0
    for i in range(len(raw)):
        if raw['head'][i] == "第" and raw['mid'][i] > 0 and raw['len'][i] < 30:
            chapnum += 1
        if chapnum >= 40 and raw['txt'][i] == "附录一：成吉思汗家族":
            chapnum = 0
        raw.loc[i, 'chap'] = chapnum
    del raw['head']
    del raw['mid']
    del raw['len']
    groupchapter = raw.groupby('chap')
    chapter = groupchapter.agg(sum)
    chapter = chapter[chapter.index != 0]
    return chapter


def Topic(chapter):
    chapters = [cut(content) for content in chapter.txt]
    dictionary = Dictionary(chapters)
    doc2bow = []
    for c in chapters:
        doc2bow.append(dictionary.doc2bow(c))
    tfidfModel = TfidfModel(doc2bow)
    tfidfVector = tfidfModel[doc2bow]
    for i in range(10, 31, 10):
        ldaModel = LdaModel(tfidfVector, id2word=dictionary, num_topics=i)
        print(f"预设主题{i}个")
        for topic in ldaModel.print_topics():
            print(topic)
    ldaModel = LdaModel(tfidfVector, id2word=dictionary, num_topics=10)
    for i in range(1, 41, 2):
        temp = chapter['txt'][i]
        tempBow = dictionary.doc2bow(cut(temp))
        tempIdf = tfidfModel[tempBow]
        topic = ldaModel.get_document_topics(tempIdf)
        resultTopic = max(topic, key=lambda x: x[1])
        print(f"第{i}章属于主题{resultTopic[0]}的概率是{resultTopic[1]}")

Topic(chapter())